import pandas as pd
import numpy as np
import matplotlib.pyplot as plt

plt.style.use('ggplot')
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False

CITIES_DISTRICTS = {
            '亳州市': '蒙城县,涡阳县,利辛县,谯城区,亳州经开区',
            '六安市': '舒城县,金寨县,霍邱县,霍山县,金安区,裕安区,六安经开区,叶集区',
            '合肥市': '肥东县,肥西县,长丰县,庐江县,巢湖市,瑶海区,庐阳区,蜀山区,包河区,合肥高新区,合肥经开区,合肥新站区,巢湖经开区',
            '安庆市': '桐城市,怀宁县,潜山县,太湖县,岳西县,宿松县,望江县,迎江区,大观区,宜秀区,安庆经开区',
            '宣城市': '宁国市,广德县,泾县,郎溪县,绩溪县,旌德县,宣州区,宣城经开区', '宿州市': '砀山县,萧县,灵璧县,泗县,埇桥区,宿州经开区',
            '池州市': '青阳县,东至县,石台县,贵池区,池州经开区', '淮北市': '濉溪县,相山区,烈山区,杜集区,淮北经开区',
            '淮南市': '凤台县,田家庵区,谢家集区,大通区,潘集区,八公山区,淮南经开区,毛集实验区,寿县',
            '滁州市': '天长市,明光市,来安县,全椒县,定远县,凤阳县,琅琊区,南谯区,滁州经开区',
            '芜湖市': '芜湖县,繁昌县,南陵县,无为县,镜湖区,鸠江区,弋江(高新）区,三山区,芜湖经开区',
            '蚌埠市': '固镇县,五河县,怀远县,龙子湖区,禹会区,蚌山区,淮上区,蚌埠高新区,蚌埠经开区',
            '铜陵市': '义安区,郊区,铜官区,铜陵经开区,枞阳县', '阜阳市': '界首市,太和县,颍上县,阜南县,临泉县,颍泉区,颍州区,颍东区,阜阳经开区',
            '马鞍山市': '当涂县,含山县,和县,花山区,雨山区,博望区,马鞍山经开区,慈湖高新区',
            '黄山市': '休宁县,祁门县,歙县,黟县,屯溪区,徽州区,黄山区,黄山经开区'
        }


class AnhuiPatents:
    def __init__(self, io):
        self.patent_kinds = ['发明', '实用新型', '外观设计']
        self.appliers = ['个人', '大专院校', '科研单位', '企业', '机关团体']
        self._all_ptt, self._effective_inv = self._get_data(io)
        self._cities_districts = CITIES_DISTRICTS
        self._city = '蚌埠市'
        self.fig = plt.figure()
        # With images, the canvas adds a big margin,
        # so I found it useful to insert fig.tight_layout(pad=0) before drawing.
        self.fig.tight_layout(pad=0)

    @property
    def city(self):
        return self._city

    @city.setter
    def city(self, city_name):
        if not city_name.endswith('市'):
            city_name = city_name + '市'
        if city_name not in self._cities_districts.keys():
            raise ValueError('设置了错误的城市名：{}'.format(city_name))
        self._city = city_name

    @property
    def districts(self) -> list:
        """return: districts of the self.city's"""
        return self._cities_districts[self.city].split(',')

    def _plot2arr(self):
        self.fig.canvas.draw()
        data = np.frombuffer(self.fig.canvas.tostring_rgb(), dtype=np.uint8)
        data = data.reshape(self.fig.canvas.get_width_height()[::-1] + (3,))
        self.fig.clf()  # 每次重新画图都清空原图
        return data

    @staticmethod
    def _get_data(io):
        """
        :return: tuple, (<DataFrame专利申请与授权汇总>, <DataFrame有效发明专利>)
        <DataFrame专利申请与授权汇总> property: self._all_ptt
            日期  所在市   地区  发明  实用新型  外观设计  个人  大专院校  科研单位   企业  机关团体    类别
        0 2012-01-01  合肥市  肥东县   1     1     3   3     0     0    2     0  专利申请
        1 2012-01-01  合肥市  肥西县  25    70    89   6    15     0  163     0  专利申请
        2 2012-01-01  合肥市  长丰县   7    49     8   0     0     0   64     0  专利申请

        <DataFrame有效发明专利>  property: self._effective_inv
           日期  所在市     地区    发明  实用新型  外观设计  个人  大专院校  科研单位    企业  机关团体
        0  2017-09  合肥市    肥东县   242   NaN   NaN  15     0     0   227     0
        1  2017-09  合肥市    肥西县  1185   NaN   NaN  42     0     0  1124    19
        2  2017-09  合肥市    长丰县   292   NaN   NaN  20     1     0   271     0

        """
        patent = pd.read_excel(io, sheet_name=[1, 2, 3])
        ptt_apply, ptt_grant, effective_inv = patent.values()
        ptt_apply['类别'] = '专利申请'
        ptt_grant['类别'] = '专利授权'
        all_ptt = pd.concat([ptt_apply, ptt_grant], ignore_index=True)
        all_ptt.fillna(0, inplace=True)
        all_ptt['日期'] = pd.to_datetime(all_ptt['日期'])
        return all_ptt, effective_inv

    def bar_prov_summary(self, ptt_appr, patent_type=None, diff=False, color='b'):
        """
        ptt_appr: one of ['发明', '实用新型', '外观设计', '个人', '大专院校', '科研单位',
                          '企业', '机关团体']
        全省各市近10年按年汇总的各项统计, 不包括今年(因今年未到年底),
        如果设置diff=True,则计算年度环比变化柱状图.
        """
        if patent_type is None:
            patent_type = '专利授权'
        current = self._all_ptt.set_index('日期', drop=True)
        current = current[current['类别'] == patent_type]
        by_year = current.groupby(lambda idx: idx.year).sum()[-11:-1]
        title = '安徽省历年{}{}量'.format(ptt_appr, patent_type)

        if diff:
            by_year = by_year.diff().dropna()
            title = '安徽省{}{}量年度环比'.format(ptt_appr, patent_type)
        by_year[ptt_appr].plot(kind='bar',
                               figsize=(7, 5),
                               title=title,
                               color=color,
                               ).axes.title.set_size(15)
        for x, y in enumerate(by_year[ptt_appr].values):
            plt.text(x=x,
                     y=y,
                     s=int(y),
                     ha='center',
                     va='bottom',
                     color='r')
        plt.grid(color='g', linestyle='--', axis='y')
        return self._plot2arr()

    def _bb_vs_top5(self, series):
        """取序列前5名, 如果包括self.city不添加, 否则加上"""
        index = list(series.sort_values()[-5:].index)
        if self.city not in index:
            index.append(self.city)
        new_s = series[index]
        else_share = 1 - new_s.sum()
        return new_s.append(pd.Series({'其它各市': else_share}))

    def _explode_cc(self, s, ratio=0.15):
        """
        '''explode current city'''
        创建pie图的爆炸系数, 为了突出显示self.city,
        ratio为突出系数
        """
        e = s.copy()
        for k in e.keys():
            if k == self.city:
                e[k] = ratio
            else:
                e[k] = 0
        return e

    def pie_prov_prop(self, ptt_appr, patent_type=None, year=None):
        """指定年份,种类的全省专利申请与授权数量占比分布饼图"""
        if patent_type is None:
            patent_type = '专利授权'
        if year is None:
            year = pd.datetime.now().year
        current_y = self._all_ptt.set_index('日期', drop=True)[str(year)]
        if current_y.empty:
            raise ValueError('缺少该年度相关数据')
        month_ended = max(current_y.index).month
        subset = current_y[current_y['类别'] == patent_type]
        share = subset.groupby('所在市')[self.patent_kinds + self.appliers].sum()
        share = share.div(share.sum(), axis=1)
        top = self._bb_vs_top5(share[ptt_appr])
        top.name = ''  # pie时不在左侧显示Y轴名称
        explotion = self._explode_cc(top, 0.1)
        top.plot.pie(title='{}年1~{}月安徽省各市{}{}数占比分布'
                           ''.format(year, month_ended, ptt_appr, patent_type),
                     figsize=(7, 7), pctdistance=0.8,  # 百分比的text离圆心的距离，0~1在圆内调节
                     labeldistance=1.05,  # 圈外标签的位置
                     autopct='%.2f%%',  # 填充百分比
                     explode=explotion,  # 爆炸效果
                     startangle=90,  # 起始角度
                     fontsize=12).axes.title.set_size(16)
        return self._plot2arr()

    def line_monthly_history(self):
        """制作self.city市专利申请及授权数量历史走势折线图"""
        curr_city = self._all_ptt[self._all_ptt['所在市'] == self.city]
        monthly_change = curr_city.groupby(['日期', '类别']).sum()
        monthly_total = monthly_change['实用新型'].add(
            monthly_change['发明'].add(monthly_change['外观设计']))
        self.fig, axes = plt.subplots(2, 2, sharex='all', sharey='all', figsize=(10, 6))
        plot_info = {(0, 0): ('全部类型', monthly_total),
                     (0, 1): ('发明', monthly_change['发明']),
                     (1, 0): ('实用新型', monthly_change['实用新型']),
                     (1, 1): ('外观设计', monthly_change['外观设计'])}
        for ax_loc, list_v in plot_info.items():
            list_v[1].unstack().plot(ax=axes[ax_loc],
                                     title=list_v[0],
                                     color=['g', 'r'],
                                     lw=0.8,
                                     fontsize=9)
            axes[ax_loc].title.set_size(8)
            axes[ax_loc].legend(fontsize=6)
        plt.subplots_adjust(wspace=0.05, hspace=0.15)
        plt.suptitle('{}专利申请及授权数量历史走势图'.format(self.city), fontsize=12, va='top')
        return self._plot2arr()

    @staticmethod
    def _rmm_avg(arr):
        _, *values, _ = sorted(arr)
        return np.mean(values)

    def line_invention(self, patent_type=None):
        """单独查看发明专利历史情况"""
        if patent_type is None:
            patent_type = '专利授权'
        curr_city = self._all_ptt[self._all_ptt['所在市'] == self.city]
        monthly_change = curr_city.groupby(['日期', '类别']).sum()
        current = monthly_change['发明'][:, patent_type]
        current.plot(style='b-',
                     lw=0.8,
                     label='历史数据曲线',
                     grid=True,
                     figsize=(10, 6),
                     title='{}发明{}历史数量'.format(self.city, patent_type),
                     fontsize=12)
        avg = self._rmm_avg(current)
        plt.axhline(y=avg,
                    xmin=0.05,
                    xmax=0.95,
                    color='r',
                    ls='--',
                    label='均值(排除最大与最小值后)',
                    lw=0.8)  # Add a horizontal line across the axis.
        plt.text(x=monthly_change.index.levels[0][1],
                 y=avg * 1.1,
                 s='历年来发明专利的{}数量:\n平均{}件/月'.format(patent_type[-2:], int(avg)),
                 fontsize=7,
                 bbox=dict(boxstyle='round, pad=0.6', fc='y', ec='k', lw=1, alpha=0.5),
                 va='bottom')
        plt.legend()
        return self._plot2arr()

    def bars_annual_share(self, patent_type=None):
        """所在市专利授权或者专利申请的各项年度占比"""
        if patent_type is None:
            patent_type = '专利授权'
        patents = self._all_ptt[self._all_ptt['类别'] == patent_type]
        total = patents.groupby(
            [patents['日期'].map(lambda tm: tm.year), '所在市']).sum()
        annual_share = total / total.sum(level='日期')
        curr_mul_index = pd.MultiIndex.from_arrays(
            [annual_share.index.levels[0], [self.city] * len(annual_share.index.levels[0])],
            names=['日期', '所在市'])
        curr_city = annual_share.reindex(curr_mul_index).reset_index().set_index('日期')

        curr_city_total = curr_city['发明'].add(curr_city['实用新型'].add(curr_city['外观设计']))
        plot_info = {
            (0, 0): ('全部专利类型', curr_city_total),
            (0, 1): ('发明{}'.format(patent_type[-2:]), curr_city['发明']),
            (1, 0): ('实用新型{}'.format(patent_type[-2:]), curr_city['实用新型']),
            (1, 1): ('外观设计{}'.format(patent_type[-2:]), curr_city['外观设计'])
        }
        self.fig, axes = plt.subplots(2, 2, sharex='all', sharey='all', figsize=(10, 5))
        for ax_loc, list_v in plot_info.items():
            list_v[1].plot(kind='bar',
                           ax=axes[ax_loc],
                           title=list_v[0],
                           fontsize=8)
            axes[ax_loc].title.set_size(10)
            for x, y in enumerate(list_v[1].values):
                axes[ax_loc].text(x=x,
                                  y=y,
                                  s=format(y, '.2%'),
                                  color='r',
                                  ha='center',
                                  va='bottom',
                                  fontsize=8)
        plt.subplots_adjust(wspace=0.1, hspace=0.15)  # 调整子图水平及垂直间距
        plt.suptitle('{}专利{}总数量年度全省占比'.format(self.city, patent_type[-2:]),
                     fontsize=12,
                     va='top')
        return self._plot2arr()

    def annual_rank(self, ptt_appr, patent_type=None):
        if patent_type is None:
            patent_type = '专利授权'
        patents = self._all_ptt[self._all_ptt['类别'] == patent_type]
        total = patents.groupby(
            [patents['日期'].map(lambda tm: tm.year), '所在市']).sum()
        annual_share = total / total.sum(level='日期')
        cities_rank = annual_share.groupby(level='日期').rank(method='dense')
        curr_mul_index = pd.MultiIndex.from_arrays(
            [annual_share.index.levels[0], [self.city] * len(annual_share.index.levels[0])],
            names=['日期', '所在市'])
        curr_city_rank = cities_rank.reindex(curr_mul_index).reset_index()
        curr_city_rank.set_index('日期', drop=True, inplace=True)
        curr_city_rank[ptt_appr].sort_index().plot(
            kind='bar',
            figsize=(10, 6),
            title='{}{}{}数量全省排名(总数16)'.format(self.city, ptt_appr, patent_type)
        ).axes.title.set_size(15)
        for x, y in enumerate(curr_city_rank[ptt_appr].values):
            plt.text(x=x,
                     y=y,
                     s='第%d名' % int(y),
                     ha='center',
                     va='bottom',
                     color='b')
        return self._plot2arr()

    def _corr_period(self, patent_type):
        curr_city = self._all_ptt[self._all_ptt['所在市'] == self.city]
        last_date = curr_city[curr_city['类别'] == patent_type]['日期'].sort_values().iloc[-1]
        last_y, last_m = last_date.year, last_date.month
        corr_period = []
        for month in range(1, last_m + 1):
            corr_period.extend(['%d-%d' % (y, month) for y in [last_y - 1, last_y]])
        corr_period = pd.to_datetime(corr_period)
        corr_period_data = curr_city.set_index('日期').loc[corr_period]
        return corr_period_data, last_date

    def bars_district_corr(self,
                           patent_type=None,
                           single=False,
                           ptt_appr=None,
                           district=None):
        if patent_type is None:
            patent_type = '专利授权'
        width = 0.35
        corr_period_data, last_date = self._corr_period(patent_type)
        last_y, last_m = last_date.year, last_date.month
        x_tickes = np.array(range(last_m))  # x轴刻度
        if not single:
            self.fig.set_size_inches(18, 30)
            n = 0
            for district in self.districts:
                by_district = corr_period_data[
                    (corr_period_data['地区'] == district) &
                    (corr_period_data['类别'] == patent_type)]
                for patent_kind in self.patent_kinds:
                    self.fig.add_subplot(len(self.districts), len(self.patent_kinds), n+1)
                    ytick_max = int(max(by_district[patent_kind]) + 1)
                    if ytick_max < 20:
                        y_tickes = list(range(0, ytick_max))
                        plt.yticks(y_tickes)
                    plt.title('{}{}{}数同期对比'
                              ''.format(district, patent_kind, patent_type), fontsize=8)
                    plt.xticks(x_tickes,  # 画x轴刻度对应的标签
                               ['{}月'.format(i) for i in range(1, last_m + 1)],
                               fontsize=8)
                    plt.bar(x_tickes,
                            by_district[str(last_y - 1)][patent_kind],
                            width=width,
                            label='{}年'.format(last_y - 1))
                    plt.bar(x_tickes + width,
                            by_district[str(last_y)][patent_kind],
                            width=width,
                            label='{}年'.format(last_y))
                    if n % 3 == 0:
                        plt.legend(fontsize=8)
                    n += 1
        else:
            if ptt_appr is None or district is None:
                raise ValueError(
                    "When single is True, ptt_appr and district can't be None"
                )
            by_district = corr_period_data[(corr_period_data['地区'] == district) &
                                           (corr_period_data['类别'] == patent_type)]
            self.fig.set_size_inches(10, 6)
            plt.title('{}{}{}数同期对比'.format(district, ptt_appr, patent_type),
                      fontsize=12)
            plt.xticks(x_tickes,  # 画x轴刻度对应的标签
                       ['{}月'.format(i) for i in range(1, last_m + 1)])
            plt.bar(x_tickes,
                    by_district[str(last_y - 1)][ptt_appr],
                    width=width,
                    label='{}年'.format(last_y - 1))
            plt.bar(x_tickes + width,
                    by_district[str(last_y)][ptt_appr],
                    width=width,
                    label='{}年'.format(last_y))
            diff_arr = (by_district[str(last_y)][ptt_appr].values -
                        by_district[str(last_y-1)][ptt_appr].values)
            for x, y in enumerate(diff_arr):
                plt.text(x=x + width/2,
                         y=by_district[str(last_y)][ptt_appr].values[x],
                         s='{}{}件'.format('↑' if y >= 0 else '↓', abs(y)),
                         ha='left',
                         va='bottom')
            plt.legend()
        return self._plot2arr()

